package com.yupi.yuaiagent.rag;

import com.yupi.yuaiagent.config.BaiduTranslateConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.regex.Pattern;

/**
 * 查询重写器 - 集成百度翻译API
 * 在查询重写前，先将英文内容翻译成中文，适配中文知识库
 */
@Slf4j
@Component
public class QueryRewriter {

    private final QueryTransformer queryTransformer;

    @Autowired
    private BaiduTranslateService baiduTranslateService;

    @Autowired
    private BaiduTranslateConfig config;

    // 英文检测正则表达式
    private static final Pattern ENGLISH_PATTERN = Pattern.compile("[a-zA-Z]");

    public QueryRewriter(ChatModel dashscopeChatModel) {
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        // 创建查询重写转换器
        queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 执行查询重写
     * 处理逻辑：
     * 1. 检测是否包含英文
     * 2. 如果包含英文且配置启用，先翻译为中文
     * 3. 对中文查询执行重写优化
     * 4. 返回优化后的查询
     *
     * @param prompt 用户输入的查询
     * @return 优化重写后的查询
     */
    public String doQueryRewrite(String prompt) {
        if (prompt == null || prompt.trim().isEmpty()) {
            return prompt;
        }

        String processedQuery = prompt;

        try {
            // 步骤1: 检测并处理英文内容
            if (containsEnglish(prompt) && config.isEnableEnToZh()) {
                log.debug("检测到英文内容，开始翻译: {}", prompt);
                String translatedText = baiduTranslateService.enToZh(prompt);

                if (translatedText != null && !translatedText.equals(prompt)) {
                    processedQuery = translatedText;
                    log.info("英文翻译完成: {} -> {}", prompt, processedQuery);
                } else {
                    log.warn("英文翻译失败或返回原文，使用原始查询: {}", prompt);
                    processedQuery = prompt;
                }
            }

            // 步骤2: 执行查询重写优化
//            log.debug("开始查询重写优化: {}", processedQuery);
//            Query query = new Query(processedQuery);
//            Query transformedQuery = queryTransformer.transform(query);
//            String rewrittenQuery = transformedQuery.text();
//
//            log.info("查询重写完成: {} -> {}", processedQuery, rewrittenQuery);
            return processedQuery;

        } catch (Exception e) {
            log.error("查询重写过程失败，返回原始查询: {}", prompt, e);
            return prompt; // 出错时返回原始查询
        }
    }

    /**
     * 检测文本是否包含英文字符
     * @param text 待检测文本
     * @return true如果包含英文，false否则
     */
    private boolean containsEnglish(String text) {
        return ENGLISH_PATTERN.matcher(text).find();
    }
}